"""
!/usr/bin/env python
-*- coding: utf-8 -*-
@CreateTime : 2025/3/3 10:13
@Author  :    AnimateX
@Contact :    animatex@163.com
@File    :    stitcher_by_optical_flow.py
@License :    Copyright © 2024 AnimateX. All rights reserved.
@Version :    stitcher_by_optical_flow_2025/3/3.0.1

-------------------------------------------------------------------------------
# @Description:

    ---------------------------------------------------------------------------
    [Update History]:
        2025/3/3:

-------------------------------------------------------------------------------
"""
import os
import cv2
import numpy as np
from utils import FisheyeUndistorter, FisheyeCrop, timer


class StitcherByOpticalFlow:
    def __init__(
        self,
        image_left_path: str = None,
        image_right_path: str = None,
        fov_degree: int = 195,
    ):
        if not os.path.exists(image_left_path):
            raise ValueError("Image left path does not exist!")
        self.image_left_path = image_left_path

        if not os.path.exists(image_right_path):
            raise ValueError("Image right path does not exist!")
        self.image_right_path = image_right_path

        if fov_degree < 0 or fov_degree > 360:
            raise ValueError("FOV degree must be between 0 and 360!")
        self.fov_degree = fov_degree

        # 初始鱼眼图的尺寸
        self.image_w = None
        self.image_h = None
        self.undistort_image_w = None

    def read_fisheye_image(self):
        """
        读取鱼眼图像并进行预处理
        """
        img_ori_l = cv2.imread(self.image_left_path)
        img_ori_r = cv2.imread(self.image_right_path)

        # Crop fisheye image
        cropper = FisheyeCrop()
        fe_img_l = cropper.crop(img_ori_l, is_first=True, en_mor=False)
        fe_img_r = cropper.crop(img_ori_r, is_first=False, en_mor=False)

        if fe_img_l.shape != fe_img_r.shape:
            raise ValueError("Left image dimensions must be equal to right image dimensions!")

        self.image_w = fe_img_l.shape[1]
        self.image_h = fe_img_l.shape[0]

        return fe_img_l, fe_img_r

    def undistort_fisheye_image(self):
        """
        对鱼眼图像进行去畸变
        """
        fe_img_l, fe_img_r = self.read_fisheye_image()

        image_w = fe_img_l.shape[1]
        image_h = fe_img_l.shape[0]
        self.undistort_image_w = int(fov_degree / 180 * image_h)

        # Undistort image
        undistorter = FisheyeUndistorter(
            self.fov_degree, image_w, image_h, undistort_image_w
        )

        undistort_img_l = undistorter.apply_undistortion(fe_img_l)
        undistort_img_r = undistorter.apply_undistortion(fe_img_r)

        # 初始化shifted_cams - 将左右图像水平拼接
        self.shifted_cams = np.vstack((undistort_img_l, undistort_img_r))

        return undistort_img_l, undistort_img_r

    def calculate_optical_flow(self):
        """
        计算光流
        """
        undistort_img_l, undistort_img_r = self.undistort_fisheye_image()

        if len(undistort_img_l.shape) == 3:
            gray_l = cv2.cvtColor(undistort_img_l, cv2.COLOR_BGR2GRAY)
        else:
            gray_l = undistort_img_l
        if len(undistort_img_r.shape) == 3:
            gray_r = cv2.cvtColor(undistort_img_r, cv2.COLOR_BGR2GRAY)
        else:
            gray_r = undistort_img_r

        """ 裁剪出重叠区域 """
        # 上下需要跳过的畸变区域（经验值）
        offset_val = int(80 / 1280 * self.image_h)

        # 畸变矫正后的图和拼接后的图像在宽上的差异
        half_diff = int((self.image_w * 2 - self.undistort_image_w) / 2)

        img_l_overlap_roi_r = gray_l[offset_val: self.image_h - offset_val, half_diff:]
        img_l_overlap_roi_l = gray_l[offset_val: self.image_h - offset_val, self.undistort_image_w - half_diff:]
        img_r_overlap_roi_l = gray_r[offset_val: self.image_h - offset_val, self.undistort_image_w - half_diff:]
        img_r_overlap_roi_r = gray_r[offset_val: self.image_h - offset_val, half_diff:]

        """ 计算光流 左边重叠区域 """
        flow_l = cv2.calcOpticalFlowFarneback(
            img_l_overlap_roi_l, img_r_overlap_roi_r, None, 0.5, 3, 15, 3, 5, 1.2, 0
        )

        """ 计算光流 右边重叠区域 """
        flow_r = cv2.calcOpticalFlowFarneback(
            img_r_overlap_roi_l, img_l_overlap_roi_r, None, 0.5, 3, 15, 3, 5, 1.2, 0
        )

        return flow_l, flow_r

    def stitch_fisheye_image(self, flow, img_l, img_r, offset_val, half_diff):
        """
        拼接鱼眼图像
        """
        # 根据光流的大小确定拼接区域的权重
        flow_magnitude = np.sqrt(flow[...,0]**2 + flow[...,1]**2)

        # 使用高斯滤波平滑光流幅度，避免局部噪声
        flow_magnitude = cv2.GaussianBlur(flow_magnitude, (15, 15), 0)

        # 归一化并反转（小的光流值表示更好的匹配）
        max_flow = np.max(flow_magnitude) if np.max(flow_magnitude) > 0 else 1.0
        weight = 1.0 - (flow_magnitude / max_flow)

        # 创建渐变融合权重 - 使用余弦渐变而不是线性渐变，效果更自然
        x = np.linspace(0, np.pi, img_l.shape[1])
        gradient_mask = 0.5 * (1 + np.cos(x))
        gradient_mask = np.tile(gradient_mask, (img_l.shape[0], 1))

        # 结合光流和渐变权重
        final_mask = weight * gradient_mask
        final_mask = np.clip(final_mask, 0, 1)

        # 扩展mask到3通道
        mask_3channel = np.stack([final_mask, final_mask, final_mask], axis=2)

        # 融合重叠区域
        blended_overlap = img_l * (1 - mask_3channel) + img_r * mask_3channel

        # 创建最终拼接图像
        result = np.zeros((self.image_h, self.image_w * 2, 3), dtype=np.uint8)

        # 左侧非重叠区域
        result[:, : half_diff] = self.shifted_cams[0: self.image_h, : half_diff]

        # 右侧非重叠区域
        result[:, half_diff + blended_overlap.shape[1]:] = self.shifted_cams[self.image_h:, half_diff + blended_overlap.shape[1]:]

        # 重叠区域使用融合结果
        result[offset_val: self.image_h - offset_val, half_diff: half_diff + blended_overlap.shape[1]] = blended_overlap

        return result


def test():
    """ [Config param] """
    fov_degree = 195
    img_left_path = r"C:\ws_rk\ws_2025\image-stitching\test\a.png"
    img_right_path = r"C:\ws_rk\ws_2025\image-stitching\test\b.png"

    img_ori_l = cv2.imread(img_left_path)
    img_ori_r = cv2.imread(img_right_path)

    # Crop fisheye image
    cropper = FisheyeCrop()
    fe_img_l = cropper.crop(img_ori_l, is_first=True, en_mor=False)
    fe_img_r = cropper.crop(img_ori_r, is_first=False, en_mor=False)

    if fe_img_l.shape != fe_img_r.shape:
        raise ValueError("Left image dimensions must be equal to right image dimensions!")

    image_w = fe_img_l.shape[1]
    image_h = fe_img_l.shape[0]
    undistort_image_w = int(fov_degree / 180 * image_h)

    # Undistort image
    undistorter = FisheyeUndistorter(fov_degree, image_w, image_h, undistort_image_w)
    undistort_img_l = undistorter.apply_undistortion(fe_img_l)
    undistort_img_r = undistorter.apply_undistortion(fe_img_r)

    gray1 = cv2.cvtColor(undistort_img_l, cv2.COLOR_BGR2GRAY)
    gray2 = cv2.cvtColor(undistort_img_r, cv2.COLOR_BGR2GRAY)

    # 计算光流
    # 使用Farneback算法计算稠密光流
    flow = cv2.calcOpticalFlowFarneback(gray1, gray2, None, 0.5, 3, 15, 3, 5, 1.2, 0)

    # 基于光流确定拼接区域
    # 这里我们假设图像已大致对齐，只需在重叠区域进行细调整
    h, w = gray1.shape
    mask = np.zeros((h, w), dtype=np.uint8)

    # 根据光流的大小确定拼接区域的权重
    # 光流越小，表示区域越匹配
    flow_magnitude = np.sqrt(flow[..., 0] ** 2 + flow[..., 1] ** 2)
    max_flow = np.max(flow_magnitude)
    weight = 1.0 - (flow_magnitude / max_flow)

    # 创建渐变融合权重
    x = np.linspace(0, 1, w)
    gradient_mask = np.tile(x, (h, 1))

    # 结合光流和渐变权重
    final_mask = weight * gradient_mask
    final_mask = np.clip(final_mask, 0, 1)

    # 扩展mask到3通道
    mask_3channel = np.stack([final_mask, final_mask, final_mask], axis=2)

    # 融合图像
    result = img1 * (1 - mask_3channel) + img2 * mask_3channel
    result = result.astype(np.uint8)

    timer.print_report()


if __name__ == '__main__':
    test()
